﻿using System;
using System.Diagnostics;
using Microsoft.Samples.Debugging.Native;
using TPOJ.Core.Utils;

namespace TPOJ.CoreWF
{
    /// <summary>
    /// The sandbox implemented using Windows Job Kernel Objects and mdbg sample.
    /// </summary>
    public class Sandbox
    {
        private IntPtr jobHandle;
        private int errorCode;
        private Process process;
        private NativePipeline pipeline = new NativePipeline();

        /// <summary>
        /// Error Code
        /// </summary>
        public int ErrorCode
        {
            get { return errorCode; }
            set { errorCode = value; }
        }
        private string errorMessage;

        /// <summary>
        /// Error Message
        /// </summary>
        public string ErrorMessage
        {
            get { return errorMessage; }
            set { errorMessage = value; }
        }

        /// <summary>
        /// Create a sandbox environment.
        /// </summary>
        public Sandbox()
        {
            SECURITY_ATTRIBUTES saJob = new SECURITY_ATTRIBUTES();
            saJob.Length = 12;
            saJob.lpSecurityDescriptor = IntPtr.Zero;
            saJob.bInheritHandle = true;
            jobHandle = NativeMethods.CreateJobObject(ref saJob, "JUDGE_JOB");
            //SetRestrictions();
        }

        /// <summary>
        /// Create a process within the sandbox.
        /// </summary>
        /// <param name="startInfo">ProcessStartInfo</param>
        /// <returns>Process created</returns>
        public Process CreateProcess(ProcessStartInfo startInfo)
        {
            process = Process.Start(startInfo);
            pipeline.Attach(process.Id);
            //AssignProcessToJob(process);
            return process;
        }

        /// <summary>
        /// Assign a process to the job kernel object of the sandbox.
        /// </summary>
        /// <param name="process">The process to be assigned</param>
        private void AssignProcessToJob(Process process)
        {
            this.process = process;
            bool isSuccess = NativeMethods.AssignProcessToJobObject(jobHandle, process.Handle);
            int v = NativeMethods.GetLastError();
            if (!isSuccess)
                throw new Exception("Error while assigning process to job.");
        }

        /// <summary>
        /// Set Sandbox Restrictions
        /// </summary>
        private void SetRestrictions()
        {
            //unsafe
            //{
            //    JOBOBJECT_EXTENDED_LIMIT_INFORMATION limits = new JOBOBJECT_EXTENDED_LIMIT_INFORMATION();
            //    if (!NativeMethods.SetInformationJobObject(this.jobHandle, JOBOBJECTINFOCLASS.JobObjectExtendedLimitInformation, limits, (uint)sizeof(JOBOBJECT_EXTENDED_LIMIT_INFORMATION)))
            //    {
            //        int b = NativeMethods.GetLastError();
            //        throw new Exception("Error while setting restrictions to job." + b.ToString());
            //    }  
            //}                            
        }

        /// <summary>
        /// Wait at most given time for a debug event happen.
        /// It would detect runtime errors.
        /// </summary>
        /// <param name="milliseconds">Time to wait in MS</param>
        /// <returns>If a debug event happened returns true,otherwise return false</returns>
        public bool WaitForDebugEvent(int milliseconds)
        {
            int numberOfProcess = 0;
            int numberOfThreads = 0;
            NativeEvent nativeEvent;
            while ((nativeEvent = pipeline.WaitForDebugEvent(milliseconds)) != null)
            {
                switch (nativeEvent.EventCode)
                {
                    case NativeDebugEventCode.CREATE_PROCESS_DEBUG_EVENT:
                        {
                            numberOfProcess++;
                            if (numberOfProcess > 1)
                            {
                                errorCode = -10;
                                errorMessage = "Restricted Usage Error. Not allowed to create child process.";
                                return false;
                            }
                            else
                                break;
                        }
                    case NativeDebugEventCode.CREATE_THREAD_DEBUG_EVENT:
                        {
                            numberOfThreads++;
                            if (numberOfThreads > 3)
                            {
                                errorCode = -11;
                                errorMessage = "Restricted Usage Error. Not allowed to create thread.";
                                return false;
                            }
                            else
                                break;
                        }
                    case NativeDebugEventCode.EXCEPTION_DEBUG_EVENT:
                        {
                            ExceptionNativeEvent e = nativeEvent as ExceptionNativeEvent;
                            ExceptionCode exception = e.ExceptionCode;
                            if (exception == ExceptionCode.STATUS_BREAKPOINT || exception == ExceptionCode.STATUS_SINGLESTEP)
                                break;
                            else
                            {
                                errorMessage = "Runtime Error. " + exception.GetExceptionMessage();
                            }
                            return false;
                        }
                    case NativeDebugEventCode.EXIT_PROCESS_DEBUG_EVENT:
                        return true;
                }
                pipeline.ContinueEvent(nativeEvent);
            }
            return false;
        }

        /// <summary>
        /// Kill the sandbox
        /// </summary>
        /// <returns>If success returns true,otherwise returns false</returns>
        public bool Kill()
        {
            bool success = NativeMethods.TerminateJobObject(jobHandle, 0);
            NativeMethods.CloseHandle(jobHandle);
            return success;
        }

    }
}
